use aes::Aes256;
use cbc::{Encryptor, cipher::{block_padding::Pkcs7, BlockEncryptMut, KeyIvInit}};
use base64::{Engine, engine::general_purpose};
use csv::Writer;

// AES-256-CBC 加密器类型别名
pub type Aes256CbcEnc = Encryptor<Aes256>;

/// 生成包含加密数据的Excel文件（以CSV格式保存）
/// 
/// # 参数
/// - output_path: 输出文件路径
/// - encrypt_columns: 需要加密的列名（如["A", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
/// - sample_data: 示例数据，格式为 Vec<Vec<String>>，第一行为表头
pub fn generate_encrypted_excel(
    output_path: &str,
    encrypt_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
    sample_data: Vec<Vec<String>>,
) -> Result<(), Box<dyn std::error::Error>> {
    if sample_data.is_empty() {
        return Err("示例数据不能为空".into());
    }

    let headers = &sample_data[0];
    
    // 找到需要加密的列索引
    let encrypt_indices: Vec<usize> = encrypt_columns.iter()
        .map(|col| headers.iter().position(|h| h == col).ok_or("加密列未找到"))
        .collect::<Result<Vec<_>, _>>()?;

    // 创建CSV写入器
    let mut wtr = Writer::from_path(output_path)?;
    
    // 写入表头
    wtr.write_record(&sample_data[0])?;

    // 处理数据行
    for row in sample_data.iter().skip(1) {
        let mut encrypted_row = row.clone();
        
        for &idx in &encrypt_indices {
            if let Some(cell_value) = encrypted_row.get_mut(idx) {
                // 加密该列的数据
                let cipher = Aes256CbcEnc::new(key.into(), iv.into());
                let data = cell_value.as_bytes();
                
                // 创建足够大的缓冲区用于加密和填充
                let mut buf = vec![0u8; data.len() + 16]; // 预留填充空间
                buf[..data.len()].copy_from_slice(data);
                
                let encrypted_data = cipher.encrypt_padded_mut::<Pkcs7>(&mut buf, data.len())
                    .map_err(|e| format!("加密失败: {:?}", e))?;
                
                // 转换为base64编码
                *cell_value = general_purpose::STANDARD.encode(encrypted_data);
            }
        }
        
        wtr.write_record(&encrypted_row)?;
    }
    
    wtr.flush()?;
    Ok(())
}

/// 生成默认的示例数据
pub fn create_sample_data() -> Vec<Vec<String>> {
    vec![
        vec!["A".to_string(), "B".to_string(), "C".to_string()], // 表头
        vec!["秘密信息1".to_string(), "公开信息1".to_string(), "机密数据1".to_string()],
        vec!["秘密信息2".to_string(), "公开信息2".to_string(), "机密数据2".to_string()],
        vec!["秘密信息3".to_string(), "公开信息3".to_string(), "机密数据3".to_string()],
        vec!["用户密码".to_string(), "用户名".to_string(), "身份证号".to_string()],
        vec!["银行卡号".to_string(), "持卡人".to_string(), "CVV码".to_string()],
    ]
}

/// 便捷方法：生成包含默认示例数据的加密Excel文件
pub fn generate_sample_encrypted_excel(
    output_path: &str,
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    let sample_data = create_sample_data();
    let encrypt_columns = ["A", "C"]; // 加密A列和C列
    
    generate_encrypted_excel(output_path, &encrypt_columns, key, iv, sample_data)?;
    
    println!("已生成加密Excel文件: {}", output_path);
    println!("加密列: {:?}", encrypt_columns);
    println!("文件包含以下数据:");
    println!("- A列(加密): 秘密信息1, 秘密信息2, 秘密信息3, 用户密码, 银行卡号");
    println!("- B列(明文): 公开信息1, 公开信息2, 公开信息3, 用户名, 持卡人");
    println!("- C列(加密): 机密数据1, 机密数据2, 机密数据3, 身份证号, CVV码");
    
    Ok(())
} 